from com.example.tree.tree_node import TreeNode


# 小偷又发现了一个新的可行窃的地区。这个地区只有一个入口，我们称之为 root 。
#  除了 root 之外，每栋房子有且只有一个“父“房子与之相连。一番侦察之后，聪明的小偷意识到“这个地方的所有房屋的排列类似于一棵二叉树”。
#  如果 两个直接相连的房子在同一天晚上被打劫 ，房屋将自动报警。
#  给定二叉树的 root 。返回 在不触动警报的情况下 ，小偷能够盗取的最高金额 。
#
#  示例 1:
# 输入: root = [3,2,3,null,3,null,1]
# 输出: 7
# 解释:小偷一晚能够盗取的最高金额 3 + 3 + 1 = 7
#
#  示例 2:
# 输入: root = [3,4,5,1,3,null,1]
# 输出: 9
# 解释:小偷一晚能够盗取的最高金额 4 + 5 = 9
class Solution:
    def rob2(self, root: TreeNode) -> int:
        """
        每个节点可选择偷或者不偷两种状态
        当前节点选择偷时，那么两个孩子节点就不能选择偷了
        当前节点选择不偷时，两个孩子节点只需要拿最多的钱出来就行(两个孩子节点偷不偷没关系)

        当前节点选择不偷：当前节点能偷到的最大钱数 = 左孩子能偷到的钱 + 右孩子能偷到的钱
        当前节点选择偷：当前节点能偷到的最大钱数 = 左孩子选择自己不偷时能得到的钱 + 右孩子选择不偷时能得到的钱 + 当前节点的钱数
        :param root:
        :return:
        """
        def dfs(node: TreeNode) -> tuple:  # 对于节点来说有偷或者不偷两种选择(res[0] 不偷该节点的最大获益， res[1] 偷该节点的最大获益)
            if not node:
                return 0, 0
            leftRes = dfs(node.left)
            rightRes = dfs(node.right)
            resNotPickRoot = max(leftRes) + max(rightRes)  # 不偷该节点就可以偷节点的左右子节点
            resPickRoot = node.val + leftRes[0] + rightRes[0]  # 偷该节点就不能偷节点的左右子节点
            return resNotPickRoot, resPickRoot
        return max(dfs(root))

    def rob1(self, root: TreeNode) -> int:
        """
        解法一(超时)：
        在于是否选择根节点
        选择了根节点(可以加上它的孙子节点)就不能再选择左右子节点
        不选择根就选择左右子节点
        两者取较大的
        :param root:
        :return:
        """
        def dfs(node: TreeNode) -> int:
            """
            暴力法(超时)
            :param node:
            :return:
            """
            if not node:
                return 0
            # 情况一：选择了根节点 + 加上它的四个孙子节点
            tmpRes1 = node.val
            tmpRes1 += dfs(node.left.left) + dfs(node.left.right) if node.left else 0
            tmpRes1 += dfs(node.right.left) + dfs(node.right.right) if node.right else 0
            # 情况二：选择左右两个子节点
            tmpRes2 = dfs(node.left) + dfs(node.right)
            return max(tmpRes1, tmpRes2)

        cachedRess = {}

        def dfsWithDict(node: TreeNode) -> int:
            """
            使用dict进行优化(超时)
            算完一个节点就将结果放到dict中
            优化：使用dict记录已经计算过的节点的值
            key 节点; value 节点对应的最大获益结果
            :param node:
            :return:
            """
            if not node:
                return 0
            if tmpRes := cachedRess.get(node):  # 如果已经计算过该节点的值，直接返回
                return tmpRes
            else:  # 没计算过该节点的值，计算完先cache，然后返回
                # 情况一：选择了根节点 + 加上它的四个孙子节点
                tmpRes1 = node.val
                tmpRes1 += dfs(node.left.left) + dfs(node.left.right) if node.left else 0
                tmpRes1 += dfs(node.right.left) + dfs(node.right.right) if node.right else 0
                # 情况二：选择左右两个子节点
                tmpRes2 = dfs(node.left) + dfs(node.right)
                cachedRess[node] = max(tmpRes1, tmpRes2)
                return max(tmpRes1, tmpRes2)

        # return dfs(root)
        return dfsWithDict(root)

    def rob(self, root: TreeNode) -> int:
        return self.rob2(root)


if __name__ == "__main__":
    """
                 3
          4             5
      1       3              1
    """
    root = TreeNode(3)
    root.left, root.right = TreeNode(4), TreeNode(5)
    root.left.left, root.left.right = TreeNode(1), TreeNode(3)
    root.right.right = TreeNode(1)
    # root = [3, 2, 3, null, 3, null, 1]
    # root = TreeNode(3)
    # root.left, root.right = TreeNode(2), TreeNode(3)
    # root.left.right = TreeNode(3)
    # root.right.right = TreeNode(1)
    print(Solution().rob(root))
